#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
数据库性能分析脚本
对比 hydrometric_data_regular (普通PostgreSQL) 和 hydrometric_data_timeseries (TimescaleDB) 的查询性能
"""

import json
import os
import time
from datetime import datetime
from typing import Dict, Tuple

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import psycopg2
import seaborn as sns
from dotenv import load_dotenv
from loguru import logger

# 加载环境变量
load_dotenv()

# 设置中文字体支持 - 针对Mac系统优化
import platform

if platform.system() == 'Darwin':  # Mac系统
    # Mac系统中文字体配置
    plt.rcParams['font.sans-serif'] = ['PingFang SC', 'Hiragino Sans GB', 'STHeiti', 'Arial Unicode MS', 'SimHei']
else:
    # 其他系统字体配置
    plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans', 'Arial Unicode MS', 'Noto Sans CJK SC']

plt.rcParams['axes.unicode_minus'] = False

def setup_chinese_font():
    """动态设置中文字体"""
    try:
        import matplotlib.font_manager as fm
        
        # 清除matplotlib字体缓存（兼容不同版本的matplotlib）
        try:
            # 新版本matplotlib
            if hasattr(fm.fontManager, 'addfont'):
                fm.fontManager.__init__()
            # 旧版本matplotlib
            elif hasattr(fm, '_rebuild'):
                fm._rebuild()
        except Exception:
            # 如果缓存清理失败，继续执行
            pass
        
        # 获取所有可用字体
        available_fonts = set([f.name for f in fm.fontManager.ttflist])
        
        # Mac系统优先字体列表 - 根据实际可用性调整顺序
        if platform.system() == 'Darwin':
            preferred_fonts = [
                'Hiragino Sans GB',  # 在Mac上最常见的中文字体
                'Arial Unicode MS',  # 备选字体
                'PingFang SC',
                'PingFang HK', 
                'PingFang TC',
                'Heiti TC',
                'Songti SC',
                'STHeiti',
                'Heiti SC',
                'Apple LiGothic Medium',
                'SimHei'
            ]
        else:
            preferred_fonts = [
                'SimHei',
                'Microsoft YaHei',
                'DejaVu Sans',
                'Arial Unicode MS',
                'Noto Sans CJK SC'
            ]
        
        # 查找第一个可用的中文字体
        found_font = None
        for font in preferred_fonts:
            if font in available_fonts:
                found_font = font
                break
        
        if found_font:
            plt.rcParams['font.sans-serif'] = [found_font]
            plt.rcParams['axes.unicode_minus'] = False
            logger.info(f"成功设置中文字体: {found_font}")
            return True
        else:
            # 如果没找到，尝试使用系统默认字体
            logger.warning("未找到预设中文字体，尝试使用系统默认字体")
            if platform.system() == 'Darwin':
                plt.rcParams['font.sans-serif'] = ['Arial Unicode MS', 'DejaVu Sans']
            else:
                plt.rcParams['font.sans-serif'] = ['DejaVu Sans']
            plt.rcParams['axes.unicode_minus'] = False
            return False
            
    except Exception as e:
        logger.error(f"字体设置失败: {e}")
        # 使用最基本的配置
        plt.rcParams['font.sans-serif'] = ['DejaVu Sans']
        plt.rcParams['axes.unicode_minus'] = False
        return False

# 设置中文字体
setup_chinese_font()


class DatabasePerformanceAnalyzer:
    def __init__(self, db_config: Dict[str, str], verbose: bool = False):
        """
        初始化数据库性能分析器
        
        Args:
            db_config: 数据库连接配置
            verbose: 控制日志详细程度
        """
        self.db_config = db_config
        self.verbose = verbose  # 控制日志详细程度
        self.results = {
            'basic_queries': [],
            'aggregation_queries': [],
            'summary': {}
        }

    def connect_db(self):
        """建立数据库连接"""
        try:
            conn = psycopg2.connect(**self.db_config)
            return conn
        except Exception as e:
            logger.error(f"数据库连接失败: {e}")
            return None

    def execute_query_with_timing(self, conn, query: str, description: str) -> Tuple[float, int]:
        """
        执行查询并记录执行时间
        
        Args:
            conn: 数据库连接
            query: SQL查询语句
            description: 查询描述
            
        Returns:
            Tuple[执行时间(秒), 结果行数]
        """
        cursor = conn.cursor()

        try:
            # 记录开始时间
            start_time = time.time()

            # 执行查询
            cursor.execute(query)
            results = cursor.fetchall()

            # 记录结束时间
            end_time = time.time()
            execution_time = end_time - start_time

            # 根据verbose设置控制日志输出详细程度
            if self.verbose or len(results) <= 1000:
                logger.info(f"{description}: {execution_time:.4f}秒, 返回{len(results):,}行")
            else:
                # 简化输出，只显示关键信息
                logger.info(f"{description}: {execution_time:.4f}秒")
            return execution_time, len(results)
        except Exception as e:
            logger.error(f"查询执行失败 - {description}: {e}")
            return 0.0, 0
        finally:
            cursor.close()

    def run_basic_queries(self):
        """执行基础查询性能测试"""
        logger.info("=== 基础查询性能测试 ===")

        conn = self.connect_db()
        if not conn:
            return

        # 测试查询列表
        queries = [
            {
                'name': '前1000行数据查询',
                'regular_sql': """
                    SELECT * FROM hydrometric_data_regular 
                    ORDER BY timestamp DESC 
                    LIMIT 1000
                """,
                'timeseries_sql': """
                    SELECT * FROM hydrometric_data_timeseries 
                    ORDER BY timestamp DESC 
                    LIMIT 1000
                """
            },
            {
                'name': '特定站点查询(ST000001-ST000005)',
                'regular_sql': """
                    SELECT * FROM hydrometric_data_regular 
                    WHERE station_id IN ('ST000001', 'ST000002', 'ST000003', 'ST000004', 'ST000005')
                    ORDER BY timestamp DESC limit 1000
                """,
                'timeseries_sql': """
                    SELECT * FROM hydrometric_data_timeseries 
                    WHERE station_id IN ('ST000001', 'ST000002', 'ST000003', 'ST000004', 'ST000005')
                    ORDER BY timestamp DESC limit 1000
                """
            },
            {
                'name': '时间范围查询(2025-08-03到2025-09-03)',
                'regular_sql': """
                    SELECT * FROM hydrometric_data_regular 
                    WHERE timestamp BETWEEN '2025-08-03' AND '2025-09-03'
                    AND station_id IN ('ST000001', 'ST000002', 'ST000003', 'ST000004', 'ST000005')
                    ORDER BY timestamp DESC limit 1000
                """,
                'timeseries_sql': """
                    SELECT * FROM hydrometric_data_timeseries 
                    WHERE timestamp BETWEEN '2025-08-03' AND '2025-09-03'
                    AND station_id IN ('ST000001', 'ST000002', 'ST000003', 'ST000004', 'ST000005')
                    ORDER BY timestamp DESC limit 1000
                """
            }
        ]

        for query_info in queries:
            logger.info(f"--- {query_info['name']} ---")

            # 执行普通表查询
            regular_time, regular_rows = self.execute_query_with_timing(
                conn, query_info['regular_sql'], f"普通表 - {query_info['name']}"
            )

            # 执行时序表查询
            timeseries_time, timeseries_rows = self.execute_query_with_timing(
                conn, query_info['timeseries_sql'], f"时序表 - {query_info['name']}"
            )

            # 记录结果
            self.results['basic_queries'].append({
                'query_name': query_info['name'],
                'regular_time': regular_time,
                'regular_rows': regular_rows,
                'timeseries_time': timeseries_time,
                'timeseries_rows': timeseries_rows,
                'performance_ratio': regular_time / timeseries_time if timeseries_time > 0 else 0
            })

        conn.close()

    def run_aggregation_queries(self):
        """执行聚合查询性能测试"""
        logger.info("=== 聚合查询性能测试 ===")

        conn = self.connect_db()
        if not conn:
            return

        # 聚合查询列表
        agg_queries = [
            {
                'name': '按日统计最大值和最小值',
                'regular_sql': """
                    SELECT 
                        station_id,
                        DATE(timestamp) as date,
                        MAX(flow_rate) as max_flow,
                        MIN(flow_rate) as min_flow,
                        COUNT(*) as record_count
                    FROM hydrometric_data_regular 
                    WHERE timestamp BETWEEN '2025-08-03' AND '2025-09-03'
                    AND station_id IN ('ST000001', 'ST000002', 'ST000003', 'ST000004', 'ST000005')
                    GROUP BY station_id, DATE(timestamp)
                    ORDER BY station_id, date
                """,
                'timeseries_sql': """
                    SELECT 
                        station_id,
                        DATE(timestamp) as date,
                        MAX(flow_rate) as max_flow,
                        MIN(flow_rate) as min_flow,
                        COUNT(*) as record_count
                    FROM hydrometric_data_timeseries 
                    WHERE timestamp BETWEEN '2025-08-03' AND '2025-09-03'
                    AND station_id IN ('ST000001', 'ST000002', 'ST000003', 'ST000004', 'ST000005')
                    GROUP BY station_id, DATE(timestamp)
                    ORDER BY station_id, date
                """
            },
            {
                'name': '按站点统计总体指标',
                'regular_sql': """
                    SELECT 
                        station_id,
                        COUNT(*) as total_records,
                        AVG(flow_rate) as avg_flow,
                        MAX(flow_rate) as max_flow,
                        MIN(flow_rate) as min_flow,
                        STDDEV(flow_rate) as stddev_flow
                    FROM hydrometric_data_regular 
                    WHERE timestamp BETWEEN '2025-08-03' AND '2025-09-03'
                    AND station_id IN ('ST000001', 'ST000002', 'ST000003', 'ST000004', 'ST000005')
                    GROUP BY station_id
                    ORDER BY station_id
                """,
                'timeseries_sql': """
                    SELECT 
                        station_id,
                        COUNT(*) as total_records,
                        AVG(flow_rate) as avg_flow,
                        MAX(flow_rate) as max_flow,
                        MIN(flow_rate) as min_flow,
                        STDDEV(flow_rate) as stddev_flow
                    FROM hydrometric_data_timeseries 
                    WHERE timestamp BETWEEN '2025-08-03' AND '2025-09-03'
                    AND station_id IN ('ST000001', 'ST000002', 'ST000003', 'ST000004', 'ST000005')
                    GROUP BY station_id
                    ORDER BY station_id
                """
            }
        ]

        for query_info in agg_queries:
            logger.info(f"--- {query_info['name']} ---")

            # 执行普通表查询
            regular_time, regular_rows = self.execute_query_with_timing(
                conn, query_info['regular_sql'], f"普通表 - {query_info['name']}"
            )

            # 执行时序表查询
            timeseries_time, timeseries_rows = self.execute_query_with_timing(
                conn, query_info['timeseries_sql'], f"时序表 - {query_info['name']}"
            )

            # 记录结果
            self.results['aggregation_queries'].append({
                'query_name': query_info['name'],
                'regular_time': regular_time,
                'regular_rows': regular_rows,
                'timeseries_time': timeseries_time,
                'timeseries_rows': timeseries_rows,
                'performance_ratio': regular_time / timeseries_time if timeseries_time > 0 else 0
            })

        conn.close()

    def generate_visualizations(self, output_dir: str):
        """生成性能对比可视化图表"""
        logger.info("=== 生成可视化图表 ===")
        
        # 确保中文字体设置正确
        setup_chinese_font()
        
        # 创建输出目录
        os.makedirs(output_dir, exist_ok=True)
        
        # 合并所有查询结果
        all_queries = self.results['basic_queries'] + self.results['aggregation_queries']
        
        if not all_queries:
            logger.warning("没有查询结果可供可视化")
            return

        # 创建DataFrame
        df = pd.DataFrame(all_queries)

        # 1. 执行时间对比柱状图
        plt.figure(figsize=(14, 8))

        x = np.arange(len(df))
        width = 0.35

        plt.bar(x - width / 2, df['regular_time'], width, label='普通PostgreSQL表', alpha=0.8, color='#ff7f0e')
        plt.bar(x + width / 2, df['timeseries_time'], width, label='TimescaleDB时序表', alpha=0.8, color='#2ca02c')

        plt.xlabel('查询类型')
        plt.ylabel('执行时间 (秒)')
        plt.title('数据库查询性能对比 - 执行时间')
        plt.xticks(x, [q[:15] + '...' if len(q) > 15 else q for q in df['query_name']], rotation=45, ha='right')
        plt.legend()
        plt.grid(True, alpha=0.3)
        plt.tight_layout()

        plt.savefig(os.path.join(output_dir, 'execution_time_comparison.png'), dpi=300, bbox_inches='tight')
        plt.close()

        # 2. 性能比率图
        plt.figure(figsize=(12, 6))

        colors = ['#d62728' if ratio > 1 else '#2ca02c' for ratio in df['performance_ratio']]
        bars = plt.bar(range(len(df)), df['performance_ratio'], color=colors, alpha=0.7)

        plt.axhline(y=1, color='black', linestyle='--', alpha=0.5, label='性能相等线')
        plt.xlabel('查询类型')
        plt.ylabel('性能比率 (普通表时间/时序表时间)')
        plt.title('TimescaleDB vs PostgreSQL 性能比率\n(>1表示TimescaleDB更快, <1表示PostgreSQL更快)')
        plt.xticks(range(len(df)), [q[:15] + '...' if len(q) > 15 else q for q in df['query_name']], rotation=45,
                   ha='right')

        # 添加数值标签
        for i, bar in enumerate(bars):
            height = bar.get_height()
            plt.text(bar.get_x() + bar.get_width() / 2., height + 0.01,
                     f'{height:.2f}x', ha='center', va='bottom')

        plt.legend()
        plt.grid(True, alpha=0.3)
        plt.tight_layout()

        plt.savefig(os.path.join(output_dir, 'performance_ratio.png'), dpi=300, bbox_inches='tight')
        plt.close()

        # 3. 热力图 - 查询类型vs数据库类型
        plt.figure(figsize=(10, 6))

        heatmap_data = []
        for query in all_queries:
            heatmap_data.append([query['regular_time'], query['timeseries_time']])

        heatmap_df = pd.DataFrame(heatmap_data,
                                  columns=['普通PostgreSQL', 'TimescaleDB'],
                                  index=[q[:20] + '...' if len(q) > 20 else q for q in df['query_name']])

        sns.heatmap(heatmap_df, annot=True, fmt='.4f', cmap='YlOrRd',
                    cbar_kws={'label': '执行时间 (秒)'})
        plt.title('查询执行时间热力图')
        plt.ylabel('查询类型')
        plt.xlabel('数据库类型')
        plt.tight_layout()

        plt.savefig(os.path.join(output_dir, 'execution_time_heatmap.png'), dpi=300, bbox_inches='tight')
        plt.close()

        logger.success(f"可视化图表已保存到: {output_dir}")

    def generate_report(self, output_dir: str):
        """生成详细的性能分析报告"""
        logger.info("=== 生成性能分析报告 ===")

        # 创建输出目录
        os.makedirs(output_dir, exist_ok=True)

        # 计算汇总统计
        all_queries = self.results['basic_queries'] + self.results['aggregation_queries']

        if not all_queries:
            logger.warning("没有查询结果可供分析")
            return

        # 计算平均性能
        avg_regular_time = np.mean([q['regular_time'] for q in all_queries])
        avg_timeseries_time = np.mean([q['timeseries_time'] for q in all_queries])
        avg_performance_ratio = np.mean([q['performance_ratio'] for q in all_queries if q['performance_ratio'] > 0])

        # 找出最佳和最差性能查询
        best_performance = max(all_queries, key=lambda x: x['performance_ratio'] if x['performance_ratio'] > 0 else 0)
        worst_performance = min(all_queries,
                                key=lambda x: x['performance_ratio'] if x['performance_ratio'] > 0 else float('inf'))

        # 生成报告内容
        report_content = f"""
# 数据库性能分析报告

## 测试概述

**测试时间**: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
**测试范围**: 
- 站点: ST000001 ~ ST000005
- 时间范围: 2025-08-03 到 2025-09-03
- 查询类型: 基础查询 + 聚合查询

## 性能汇总

### 平均执行时间
- **普通PostgreSQL表**: {avg_regular_time:.4f} 秒
- **TimescaleDB时序表**: {avg_timeseries_time:.4f} 秒
- **平均性能比率**: {avg_performance_ratio:.2f}x (TimescaleDB相对优势)

### 性能极值
- **最佳性能查询**: {best_performance['query_name']} (性能比率: {best_performance['performance_ratio']:.2f}x)
- **最差性能查询**: {worst_performance['query_name']} (性能比率: {worst_performance['performance_ratio']:.2f}x)

## 详细查询结果

### 基础查询性能

"""

        # 添加基础查询详情
        for i, query in enumerate(self.results['basic_queries'], 1):
            report_content += f"""
#### {i}. {query['query_name']}
- **普通表执行时间**: {query['regular_time']:.4f} 秒 (返回 {query['regular_rows']} 行)
- **时序表执行时间**: {query['timeseries_time']:.4f} 秒 (返回 {query['timeseries_rows']} 行)
- **性能比率**: {query['performance_ratio']:.2f}x
- **性能提升**: {((query['performance_ratio'] - 1) * 100):.1f}%

"""

        report_content += "\n### 聚合查询性能\n\n"

        # 添加聚合查询详情
        for i, query in enumerate(self.results['aggregation_queries'], 1):
            report_content += f"""
#### {i}. {query['query_name']}
- **普通表执行时间**: {query['regular_time']:.4f} 秒 (返回 {query['regular_rows']} 行)
- **时序表执行时间**: {query['timeseries_time']:.4f} 秒 (返回 {query['timeseries_rows']} 行)
- **性能比率**: {query['performance_ratio']:.2f}x
- **性能提升**: {((query['performance_ratio'] - 1) * 100):.1f}%

"""

        # 添加结论和建议
        report_content += f"""
## 分析结论

### 性能优势分析

1. **TimescaleDB在时序查询方面的优势**:
   - 平均性能提升: {((avg_performance_ratio - 1) * 100):.1f}%
   - 特别适合时间范围查询和聚合操作
   - 压缩功能有效减少存储空间和I/O开销

2. **查询类型适用性**:
   - 基础查询: {'TimescaleDB更优' if avg_performance_ratio > 1 else 'PostgreSQL更优'}
   - 聚合查询: {'TimescaleDB更优' if len([q for q in self.results['aggregation_queries'] if q['performance_ratio'] > 1]) > len(self.results['aggregation_queries']) / 2 else 'PostgreSQL更优'}

### 建议

1. **数据存储策略**:
   - 对于时序数据，建议使用TimescaleDB以获得更好的查询性能
   - 对于非时序数据，普通PostgreSQL表已足够

2. **查询优化**:
   - 充分利用TimescaleDB的时间分区特性
   - 合理设置压缩策略以平衡查询性能和存储成本

3. **系统架构**:
   - 混合使用两种表结构，根据数据特性选择合适的存储方式
   - 定期监控查询性能，根据实际使用情况调整优化策略

## 技术细节

### 测试环境
- PostgreSQL + TimescaleDB扩展
- 表结构参考: /Users/gaoyu/source_code/HydroPulse-DX/sql/docker-init.sql#L4-51
- 测试数据: 水文监测站点数据

### 优化配置
- TimescaleDB chunk_time_interval: 7天
- 压缩策略: 30天后自动压缩
- 分区策略: 按station_id分区，4个分区

---

*报告生成时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}*
"""

        # 保存报告
        report_path = os.path.join(output_dir, 'performance_analysis_report.md')
        with open(report_path, 'w', encoding='utf-8') as f:
            f.write(report_content)

        # 保存JSON格式的原始数据
        json_path = os.path.join(output_dir, 'performance_data.json')
        with open(json_path, 'w', encoding='utf-8') as f:
            json.dump(self.results, f, ensure_ascii=False, indent=2)

        logger.success(f"性能分析报告已保存到: {report_path}")
        logger.success(f"原始数据已保存到: {json_path}")

    def run_full_analysis(self, output_dir: str):
        """运行完整的性能分析"""
        logger.info("开始数据库性能分析...")
        logger.info(f"输出目录: {output_dir}")

        # 执行基础查询测试
        self.run_basic_queries()

        # 执行聚合查询测试
        self.run_aggregation_queries()

        # 生成可视化图表
        self.generate_visualizations(output_dir)

        # 生成分析报告
        self.generate_report(output_dir)

        logger.success("性能分析完成!")


def main():
    """主函数"""
    # 数据库连接配置 - 从环境变量读取
    db_config = {
        'host': os.getenv('DB_HOST', 'localhost'),
        'port': int(os.getenv('DB_PORT', 5432)),
        'database': os.getenv('DB_NAME', 'hydropulse'),
        'user': os.getenv('DB_USER', 'postgres'),
        'password': os.getenv('DB_PASSWORD', 'password')
    }

    # 输出目录 - 使用相对路径
    output_dir = os.path.join(os.path.dirname(__file__))

    # 创建分析器实例 (默认简洁模式)
    verbose_mode = os.getenv('PERFORMANCE_VERBOSE', 'false').lower() == 'true'
    analyzer = DatabasePerformanceAnalyzer(db_config, verbose=verbose_mode)

    # 运行完整分析
    analyzer.run_full_analysis(output_dir)


if __name__ == '__main__':
    main()
